﻿/*
  OtpKeyProv Plugin
  Copyright (C) 2011-2012 Dominik Reichl <dominik.reichl@t-online.de>

  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 2 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program; if not, write to the Free Software
  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
*/

using System;
using System.Collections.Generic;
using System.Text;
using System.IO;
using System.Security.Cryptography;
using System.Diagnostics;

using KeePassLib.Cryptography;
using KeePassLib.Cryptography.Cipher;
using KeePassLib.Cryptography.KeyDerivation;
using KeePassLib.Keys;
using KeePassLib.Utility;

namespace OtpKeyProv
{
  public enum OtpDataFmt
  {
    Hex = 0,
    Base64 = 1,
    Base32 = 4,
    Utf8 = 2,
    Dec = 3
  }

  public static class OtpUtil
  {
    public static byte[] KeyFromOtps(string[] vOtps, int iOtpsOffset,
        int iOtpsCount, byte[] pbTrfKey32, ulong uTrfRounds)
    {
      StringBuilder sb = new StringBuilder();
      for (int i = iOtpsOffset; i < (iOtpsOffset + iOtpsCount); ++i)
      {
        if (sb.Length > 0) sb.Append(':');
        sb.Append(vOtps[i].Trim());
      }

      string strKey = sb.ToString();
      byte[] pb = StrUtil.Utf8.GetBytes(strKey);
      if (pb.Length == 0) return null;

      byte[] pbKey = HashAndTransform(pb, pbTrfKey32, uTrfRounds);
      if (pbKey == null) throw new InvalidOperationException();

      return pbKey;
    }

    public static string EncryptData(byte[] pbData, byte[] pbKey32,
        byte[] pbIV16)
    {
      byte[] pbIV8 = new byte[8];
      Array.Copy(pbIV16, 0, pbIV8, 0, 8);

      byte[] pbEnc = new byte[pbData.Length];
      Array.Copy(pbData, pbEnc, pbData.Length);

      Salsa20Cipher enc = new Salsa20Cipher(pbKey32, pbIV8);
      enc.Encrypt(pbEnc, 0, pbEnc.Length);

      return ("s20://" + Convert.ToBase64String(pbEnc,
          Base64FormattingOptions.None));
    }

    public static byte[] DecryptData(string strData, byte[] pbKey32,
        byte[] pbIV16)
    {
      if (!strData.StartsWith("s20://")) return null;

      string strEnc = strData.Substring(6);
      byte[] pb = Convert.FromBase64String(strEnc);

      byte[] pbIV8 = new byte[8];
      Array.Copy(pbIV16, 0, pbIV8, 0, 8);

      Salsa20Cipher dec = new Salsa20Cipher(pbKey32, pbIV8);
      dec.Encrypt(pb, 0, pb.Length);

      return pb;
    }

    private static byte[] HashAndTransform(byte[] pbData, byte[] pbTrfKey32,
        ulong uTrfRounds)
    {
      SHA256Managed sha256 = new SHA256Managed();
      byte[] pbHash = sha256.ComputeHash(pbData);
      sha256.Clear();

      if (!AesKdf.TransformKeyManaged(pbHash, pbTrfKey32, uTrfRounds))
        return null;

      sha256 = new SHA256Managed();
      pbHash = sha256.ComputeHash(pbHash);
      sha256.Clear();

      return pbHash;
    }

    public static OtpEncryptedData EncryptSecret(byte[] pbSecret, string[] vOtps,
        int iOtpsOffset, int iOtpsCount)
    {
      OtpEncryptedData d = new OtpEncryptedData();
      CryptoRandom r = CryptoRandom.Instance;

      byte[] pbIV16 = r.GetRandomBytes(16);
      d.IV = Convert.ToBase64String(pbIV16, Base64FormattingOptions.None);

      byte[] pbTrfKey32 = r.GetRandomBytes(32);
      d.TransformationKey = Convert.ToBase64String(pbTrfKey32, Base64FormattingOptions.None);

      byte[] pbKey32 = OtpUtil.KeyFromOtps(vOtps, iOtpsOffset, iOtpsCount,
          pbTrfKey32, d.TransformationRounds);

      d.CipherText = OtpUtil.EncryptData(pbSecret, pbKey32, pbIV16);

      byte[] pbHashTrfKey32 = r.GetRandomBytes(32);
      d.PlainTextHashTransformationKey = Convert.ToBase64String(pbHashTrfKey32,
          Base64FormattingOptions.None);

      byte[] pbHash = HashAndTransform(pbSecret, pbHashTrfKey32,
          d.PlainTextHashTransformationRounds);
      d.PlainTextHash = Convert.ToBase64String(pbHash, Base64FormattingOptions.None);

      return d;
    }

    public static byte[] DecryptSecret(OtpEncryptedData d, string[] vOtps,
        int iOtpsOffset, int iOtpsCount)
    {
      try { return DecryptSecretPriv(d, vOtps, iOtpsOffset, iOtpsCount); }
      catch (Exception) { Debug.Assert(false); }
      return null;
    }

    private static byte[] DecryptSecretPriv(OtpEncryptedData d, string[] vOtps,
        int iOtpsOffset, int iOtpsCount)
    {
      if (d == null) throw new ArgumentNullException("d");

      byte[] pbTrfKey32 = Convert.FromBase64String(d.TransformationKey);
      byte[] pbKey32 = OtpUtil.KeyFromOtps(vOtps, iOtpsOffset, iOtpsCount,
          pbTrfKey32, d.TransformationRounds);
      byte[] pbIV = Convert.FromBase64String(d.IV);

      byte[] pbSecret = OtpUtil.DecryptData(d.CipherText, pbKey32, pbIV);

      byte[] pbHashTrfKey32 = Convert.FromBase64String(d.PlainTextHashTransformationKey);
      byte[] pbHash = HashAndTransform(pbSecret, pbHashTrfKey32,
          d.PlainTextHashTransformationRounds);

      if (!MemUtil.ArraysEqual(pbHash, Convert.FromBase64String(d.PlainTextHash)))
        return null;

      return pbSecret;
    }
  }
}
